from PIL import ImageOps
from PIL import Image, ImageDraw
import PIL
from diffusers.utils import load_image


def process_img_1024(imgpath,
                     img_pil=None,
                     target_shape = (1024,1024),
                     pad_color=(255,255,255)) -> PIL.Image:
    image = load_image(imgpath) if img_pil is None else img_pil # PIL
    image,_,_,_,_ = pad_image(image.convert('RGB'),
                              height=target_shape[1],
                              width=target_shape[0],
                              fill=pad_color)
    image = image.resize(target_shape)
    return image

def resize_with_aspect(image_path, img_pil=None , max_size=1024):
    """
    保持宽高比缩放图片，使最长边不超过指定尺寸
    
    :param image_path: 输入图片路径
    :param output_path: 输出图片路径
    :param max_size: 最长边限制（默认1024像素）
    """
    # with Image.open(image_path) as img:
    img = Image.open(image_path) if img_pil is None else img_pil
    # 获取原始尺寸
    width, height = img.size
    
    # 计算新尺寸（保持宽高比）
    if width > height:
        # 横版图片：宽度限制为max_size
        new_width = max_size
        new_height = int(height * (max_size / width))
    else:
        # 竖版图片：高度限制为max_size
        new_height = max_size
        new_width = int(width * (max_size / height))
    
    # 使用高质量缩略图方法
    img.thumbnail((new_width, new_height), Image.Resampling.LANCZOS)
    
    return img        


def horizontal_concat_images(images, gap=10, line_width=2, line_color='gray'):
    """
    横向拼接多张图片，并在图片之间添加分隔线
    
    Args:
        images: 图片列表 (PIL.Image 对象列表)
        gap: 图片间距 (像素)
        line_width: 分隔线宽度 (像素)
        line_color: 分隔线颜色
        
    Returns:
        拼接后的新图片 (PIL.Image 对象)
    """
    # 计算总宽度 (所有图片宽度之和 + 间距)
    widths, heights = zip(*(img.size for img in images))
    total_width = sum(widths) + gap * (len(images) - 1)
    max_height = max(heights)
    
    # 创建新图片 (白色背景)
    new_img = Image.new('RGB', (total_width, max_height), color='white')
    
    # 粘贴图片并绘制分隔线
    x_offset = 0
    for i, img in enumerate(images):
        new_img.paste(img, (x_offset, 0))
        
        # 如果不是最后一张图片，绘制分隔线
        if i < len(images) - 1:
            line_x = x_offset + img.width + gap // 2
            draw = ImageDraw.Draw(new_img)
            draw.line(
                [(line_x, 0), (line_x, max_height)],
                fill=line_color,
                width=line_width
            )
        
        x_offset += img.width + gap
    
    return new_img

def vertical_concat_images(images, gap=10, line_width=2, line_color='gray'):
    """
    垂直拼接多张图片，并在图片之间添加分隔线
    
    Args:
        images: 图片列表 (PIL.Image 对象列表)
        gap: 图片间距 (像素)
        line_width: 分隔线宽度 (像素)
        line_color: 分隔线颜色
        
    Returns:
        拼接后的新图片 (PIL.Image 对象)
    """
    # 计算总高度 (所有图片高度之和 + 间距)
    widths, heights = zip(*(img.size for img in images))
    max_width = max(widths)
    total_height = sum(heights) + gap * (len(images) - 1)
    
    # 创建新图片 (白色背景)
    new_img = Image.new('RGB', (max_width, total_height), color='white')
    
    # 粘贴图片并绘制分隔线
    y_offset = 0
    for i, img in enumerate(images):
        # 水平居中
        x_offset = (max_width - img.width) // 2
        new_img.paste(img, (x_offset, y_offset))
        
        # 如果不是最后一张图片，绘制分隔线
        if i < len(images) - 1:
            line_y = y_offset + img.height + gap // 2
            draw = ImageDraw.Draw(new_img)
            draw.line(
                [(0, line_y), (max_width, line_y)],
                fill=line_color,
                width=line_width
            )
        
        y_offset += img.height + gap
    
    return new_img

def pad_image(image, height=1024, width=1024, fill=(255,255,255)):
        w, h = image.size
        left, top, right, bottom=0,0,0,0

        if w / h < (width / height):
            right = int((width / height) * h - w)
            left = right // 2
            right = right - left
            expanded_img = ImageOps.expand(
                image, border=(left, 0, right, 0), fill=fill
            )
        else:
            bottom = int((height/width) * w - h)
            top = bottom // 2
            bottom = bottom - top
            expanded_img = ImageOps.expand(
                image, border=(0, top, 0, bottom), fill=fill
            )
        # return expanded_img, right, bottom, 0, 0
        expanded_img = expanded_img.resize((width,height))
        return expanded_img, left, top, right, bottom


def concat_half(img1_path, img2_path,concat_type = 'h'):
    # 打开两张图片
    img1 = process_img_1024(img1_path)
    img2 = process_img_1024(img2_path)
    
    # 确保两张图片高度相同
    if img1.height != img2.height:
        img2 = img2.resize((img2.width, img1.height))
    
    if concat_type == 'h':
        # 各取一半宽度
        half_width = img1.width // 2
        left_part = img1.crop((0, 0, half_width, img1.height))
        right_part = img2.crop((img2.width - half_width, 0, img2.width, img2.height))
        
        # 创建新图并拼接
        new_img = Image.new('RGB', (img1.width, img1.height))
        new_img.paste(left_part, (0, 0))
        new_img.paste(right_part, (half_width, 0))

    if concat_type == 'v':
        # 各取一半高度
        half_height = img1.height // 2
        top_part = img1.crop((0, 0, img1.width, half_height))
        bottom_part = img2.crop((0, img2.height - half_height, img2.width, img2.height))
        
        # 创建新图并拼接
        new_img = Image.new('RGB', (img1.width, img1.height))
        new_img.paste(top_part, (0, 0))
        new_img.paste(bottom_part, (0, half_height))
    if concat_type == 'c':    
        # 创建圆形遮罩
        mask = Image.new('L', img1.size, 0)
        draw = ImageDraw.Draw(mask)
        
        # 计算圆形参数（取图像短边的40%作为内圆半径）
        radius = min(img1.size) * 0.4
        center = (img1.width//2, img1.height//2)
        draw.ellipse((center[0]-radius, center[1]-radius, 
                    center[0]+radius, center[1]+radius), fill=255)
        
        # 应用遮罩
        img2.putalpha(mask)
        img1.paste(img2, (0, 0), img2)
        new_img = img1

    # new_img.save(output_path)
    return new_img
